#include "AttachDetachFilter.h"

#include "Init.h"
#include "Utils.h"

BOOLEAN IsDeviceEqualTo(PDEVICE_OBJECT pDeviceObject, const wchar_t* pwsDeviceName);

#pragma alloc_text(PAGE, AttachToFSDeviceObject)

NTSTATUS AttachToFSDeviceObject(IN const PDEVICE_OBJECT pTargetCDO)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    PDEVICE_OBJECT pNewFilterDeviceObject = NULL;
    PFILE_HIDING_FILTER_DEVICE_EXTENSION pDeviceExt = NULL;

    PAGED_CODE();

    if (!IS_DESIRED_DEVICE_TYPE(pTargetCDO))
    {
        return STATUS_NOT_SUPPORTED;
    }

    // check for FS recognizer and skip if match
    if (IsDeviceEqualTo(pTargetCDO, L"\\FileSystem\\Fs_Rec"))
    {
        return STATUS_SUCCESS;
    }

    // create new device for attaching to stack
    ntStatus = CreateFilterDevice(pTargetCDO->DeviceType, &pNewFilterDeviceObject);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    pDeviceExt = (PFILE_HIDING_FILTER_DEVICE_EXTENSION)pNewFilterDeviceObject->DeviceExtension;

    // copy flags from pTargetCDO to new device
    COPY_FILTER_DEVICE_FLAGS(pTargetCDO, pNewFilterDeviceObject);

    // attach new device to device stack
    ntStatus = IoAttachDeviceToDeviceStackSafe(pNewFilterDeviceObject, pTargetCDO, &pDeviceExt->pNextLowerDevice);

    if (!NT_SUCCESS(ntStatus))
    {
        IoDeleteDevice(pNewFilterDeviceObject);

        return ntStatus;
    }

    ClearFlag(pNewFilterDeviceObject->Flags, DO_DEVICE_INITIALIZING);

    return STATUS_SUCCESS;
}

#pragma alloc_text(PAGE, IsDeviceEqualTo)

BOOLEAN IsDeviceEqualTo(PDEVICE_OBJECT pDeviceObject, const wchar_t* pwsDeviceNameForCheck)
{
    UNICODE_STRING usDeviceNameForCheck = {0};
    UNICODE_STRING usDeviceName = {0};
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;

    PAGED_CODE();

    RtlInitUnicodeString(&usDeviceNameForCheck, pwsDeviceNameForCheck);

    ntStatus = GetObjectName(pDeviceObject->DriverObject, &usDeviceName, PagedPool);

    if (NT_SUCCESS(ntStatus))
    {
        if (RtlCompareUnicodeString(&usDeviceName, &usDeviceNameForCheck, FALSE) == 0)
        {
            FreeUnicodeString(&usDeviceName);

            return TRUE;
        }

        FreeUnicodeString(&usDeviceName);
    }

    return FALSE;
}

void CheckVolumeAndAttachIfNeeds(IN const PDEVICE_OBJECT pVolumeDevice, IN PVOID pContext);

#pragma alloc_text(PAGE, EnumerateFSVolumesAndAttach)

NTSTATUS EnumerateFSVolumesAndAttach(IN const PDEVICE_OBJECT pTargetCDO)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;

    PAGED_CODE();
    
    ntStatus = EnumerateFSVolumes(pTargetCDO, CheckVolumeAndAttachIfNeeds, pTargetCDO);

    return ntStatus;
}

BOOLEAN IsAttachedToDevice(IN const PDEVICE_OBJECT pVolumeDevice, OUT OPTIONAL PDEVICE_OBJECT* ppAttachedDevice);

#pragma alloc_text(PAGE, CheckVolumeAndAttachIfNeeds)

void CheckVolumeAndAttachIfNeeds(IN const PDEVICE_OBJECT pVolumeDevice, IN PVOID pContext)
{
    UNICODE_STRING usBaseDeviceObjectName = {0};
    PDEVICE_OBJECT pBaseDeviceObject = NULL;
    PDEVICE_OBJECT pDiskDeviceObject = NULL;
    PDEVICE_OBJECT pNewFilterDeviceObject = NULL;
    PDEVICE_OBJECT pFileHidingFilterCDO = NULL;
    PDEVICE_OBJECT pTargetCDO = (PDEVICE_OBJECT)pContext;
    PFILE_HIDING_FILTER_DEVICE_EXTENSION pDeviceExt = NULL;
    unsigned long ulAttachTryingCounter = 0;
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;

    UNREFERENCED_PARAMETER(pContext);

    PAGED_CODE();

    // check if pVolumeDevice is CDO, has incompatible type or we already attached to this device
    if (pVolumeDevice == pTargetCDO ||
        pVolumeDevice->DeviceType != pTargetCDO->DeviceType ||
        IsAttachedToDevice(pVolumeDevice, NULL) )
    {
        return;
    }

    pBaseDeviceObject = IoGetDeviceAttachmentBaseRef(pVolumeDevice);

    ntStatus = GetObjectName(pBaseDeviceObject, &usBaseDeviceObjectName, PagedPool);

    ObDereferenceObject(pBaseDeviceObject);

    if (!NT_SUCCESS(ntStatus) || 0 < usBaseDeviceObjectName.Length)
    {
        FreeUnicodeString(&usBaseDeviceObjectName);

        return;
    }

    FreeUnicodeString(&usBaseDeviceObjectName);

    ntStatus = IoGetDiskDeviceObject(pVolumeDevice, &pDiskDeviceObject);

    if (!NT_SUCCESS(ntStatus))
    {
        return;
    }

    // check and skip shadow copy volume
    if (FILE_DEVICE_DISK == pDiskDeviceObject->DeviceType && IsDeviceEqualTo(pDiskDeviceObject, L"\\Driver\\VolSnap"))
    {
        ObDereferenceObject(pDiskDeviceObject);

        return;
    }

    ObDereferenceObject(pDiskDeviceObject);

    // create new device for attaching to stack
    ntStatus = CreateFilterDevice(pVolumeDevice->DeviceType, &pNewFilterDeviceObject);

    if (!NT_SUCCESS(ntStatus))
    {
        return;
    }

    // copy flags from pTargetCDO to new device
    COPY_FILTER_DEVICE_FLAGS(pVolumeDevice, pNewFilterDeviceObject);

    pDeviceExt = (PFILE_HIDING_FILTER_DEVICE_EXTENSION)pNewFilterDeviceObject->DeviceExtension;

    // lock CDO and attach to stack
    pFileHidingFilterCDO = GetFilterCDOAndLock();

    if (IsAttachedToDevice(pVolumeDevice, NULL))
    {
        ReleaseFilterCDOAndUnlock(&pFileHidingFilterCDO);

        return;
    }

    /*
    Microsoft: it is possible for this attachment request to fail because this device
    object has not finished initializing. This can occur if this filter loaded just as this volume was being mounted.
    */

    for (ulAttachTryingCounter = 0; ulAttachTryingCounter < 8; ++ulAttachTryingCounter)
    {
        LARGE_INTEGER liInterval = {0};

        // attach new device to device stack
        ntStatus = IoAttachDeviceToDeviceStackSafe(pNewFilterDeviceObject, pVolumeDevice, &pDeviceExt->pNextLowerDevice);

        // delay this thread if failed and give another thread chance to complete mounting
        if (!NT_SUCCESS(ntStatus))
        {
            liInterval.QuadPart = -10000 * 500; // wait 500ms

            KeDelayExecutionThread(KernelMode, FALSE, &liInterval);

            // try again
            continue;
        }

        // all success, so we can break this cycle
        ClearFlag(pNewFilterDeviceObject->Flags, DO_DEVICE_INITIALIZING);

        break;
    }

    if (!NT_SUCCESS(ntStatus))
    {
        IoDeleteDevice(pNewFilterDeviceObject);
    }

    ReleaseFilterCDOAndUnlock(&pFileHidingFilterCDO);
}

#pragma alloc_text(PAGE, IsAttachedToDevice)

BOOLEAN IsAttachedToDevice(IN const PDEVICE_OBJECT pVolumeDevice, OUT OPTIONAL PDEVICE_OBJECT* ppAttachedDevice)
{
    PDEVICE_OBJECT pCurrentDeviceObject = NULL;
    PDEVICE_OBJECT pNextDeviceObject = NULL;

    PAGED_CODE();

    pCurrentDeviceObject = IoGetAttachedDeviceReference(pVolumeDevice);

    while (NULL != pCurrentDeviceObject)
    {
        // check if our filter device object
        if (IS_FILE_HIDING_FILTER_OBJECT(pCurrentDeviceObject))
        {
            if (ARGUMENT_PRESENT(ppAttachedDevice))
            {
                *ppAttachedDevice = pCurrentDeviceObject;
            }
            else
            {
                ObDereferenceObject(pCurrentDeviceObject);
            }

            return TRUE;
        }

        pNextDeviceObject = IoGetLowerDeviceObject(pCurrentDeviceObject);

        ObDereferenceObject(pCurrentDeviceObject);

        pCurrentDeviceObject = pNextDeviceObject;
    }

    if (ARGUMENT_PRESENT(ppAttachedDevice))
    {

        *ppAttachedDevice = NULL;
    }

    return FALSE;
}

#pragma alloc_text(PAGE, DetachFilter)

void DetachFilter(IN const PDEVICE_OBJECT pFSDeviceObject)
{
    PDEVICE_OBJECT pFilterAttachedDevice = pFSDeviceObject->AttachedDevice;
    PDEVICE_OBJECT pPreviousAttachedDevice = pFSDeviceObject;

    PAGED_CODE();

    while (NULL != pFilterAttachedDevice)
    {
        if (IS_FILE_HIDING_FILTER_OBJECT(pFilterAttachedDevice))
        {
            IoDetachDevice(pPreviousAttachedDevice);
            IoDeleteDevice(pFilterAttachedDevice);

            return;
        }

        pPreviousAttachedDevice = pFilterAttachedDevice;
        pFilterAttachedDevice = pFilterAttachedDevice->AttachedDevice;
    }
}